# py/sdk.py
import qianfan
import os
from typing import List, Generator

class QianfanWrapper:
    """支持流式输出的千帆调用封装"""
    
    def __init__(self):
        self._init_qianfan()
        self.model = "ERNIE-Speed-8K"  # 使用速度优化模型
        self.max_retries = 2  # 减少重试次数
    
    def _init_qianfan(self):
        os.environ["QIANFAN_ACCESS_KEY"] = os.getenv("QIANFAN_AK", "ALTAKiwnHEggAvYXbrd21K3luc")
        os.environ["QIANFAN_SECRET_KEY"] = os.getenv("QIANFAN_SK", "0d9cd08f6a924dac8673281d172ba485")
    
    def query_stream(self, prompt: str, timeout: int = 8) -> Generator[str, None, None]:
        """流式查询方法"""
        chat = qianfan.ChatCompletion(model=self.model)
        try:
            resp = chat.do(
                messages=[{"role": "user", "content": prompt}],
                stream=True,
                temperature=0.3,  # 降低随机性加速响应
                request_timeout=timeout,
                incremental_output=True  # 关键参数：立即流式输出
            )
            
            for chunk in resp:
                if 'result' in chunk:
                    yield chunk['result']
                    
        except Exception as e:
            yield f"API调用失败: {str(e)}"

def complete_answer(cues: List[str]) -> Generator[str, None, None]:
    # 明确表示只返回生成器
    wrapper = QianfanWrapper()
    prompt = f"""
    【问题】{cues[0]}
    并提供了这些关系作为回答线索{', '.join(cues[1:]) if len(cues) > 1 else '无'}
    请按照以下方式回答：
    1. 对问题展开分析并回答
    2. 结合提供的关系回答问题
    3. 重点揭示这些关系背后的历史关联和智慧
    """
    try:
        # 直接返回生成器（不处理strip）
        yield from wrapper.query_stream(prompt, timeout=10)
    except Exception as e:
        yield f"生成失败：{str(e)}"